package org.lql.chaos.rocketmq.convert;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.MessageLite;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.converter.MessageConversionException;
import org.springframework.messaging.converter.SmartMessageConverter;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.util.Assert;

import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: liangqinglong
 * @date: 2025-08-15 17:40
 * @description: 支持多种数据类型的转换器，可处理String、byte[]、Protobuf和JSON类型
 **/
public class MultiTypeMessageConverter implements SmartMessageConverter {
    
    private static final Logger log = LoggerFactory.getLogger(MultiTypeMessageConverter.class);
    
    // 内容类型常量
    private static final String CONTENT_TYPE_TEXT_PLAIN = "text/plain";
    private static final String CONTENT_TYPE_OCTET_STREAM = "application/octet-stream";
    private static final String CONTENT_TYPE_PROTOBUF = "application/x-protobuf";
    private static final String CONTENT_TYPE_JSON = "application/json";
    
    // Protobuf方法缓存，提高性能
    private final ConcurrentHashMap<Class<?>, Method> protobufParseMethodCache = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<Class<?>, Method> protobufToByteArrayMethodCache = new ConcurrentHashMap<>();
    
    private final ObjectMapper objectMapper;

    public MultiTypeMessageConverter(ObjectMapper objectMapper) {
        Assert.notNull(objectMapper, "ObjectMapper must not be null");
        this.objectMapper = objectMapper;
    }

    @Override
    public Object fromMessage(Message<?> message, Class<?> targetClass) {
        return fromMessage(message, targetClass, null);
    }

    @Override
    public Object fromMessage(Message<?> message, Class<?> targetClass, Object conversionHint)
            throws MessageConversionException {
        byte[] body = (byte[]) message.getPayload();
        try {
            if (targetClass == String.class) {
                return new String(body, StandardCharsets.UTF_8);
            }
            if (targetClass == byte[].class) {
                return body;
            }
            // Protobuf
            if (MessageLite.class.isAssignableFrom(targetClass)) {
                return parseProtobufMessage(body, targetClass);
            }
            // JSON
            return objectMapper.readValue(body, targetClass);
        } catch (Exception e) {
            log.error("Failed to convert message to target class: {}", targetClass.getName(), e);
            throw new MessageConversionException("Failed to convert message to " + targetClass.getName(), e);
        }
    }

    /**
     * 解析Protobuf消息，使用缓存提高性能
     * 
     * @param body 消息体
     * @param targetClass 目标类
     * @return 解析后的对象
     * @throws NoSuchMethodException 找不到方法
     * @throws IllegalAccessException 方法访问异常
     * @throws InvocationTargetException 方法调用异常
     * @throws InvalidProtocolBufferException Protobuf解析异常
     */
    private Object parseProtobufMessage(byte[] body, Class<?> targetClass) 
            throws NoSuchMethodException, IllegalAccessException, InvocationTargetException, InvalidProtocolBufferException {
        Method parseMethod = protobufParseMethodCache.computeIfAbsent(targetClass, cls -> {
            try {
                return cls.getMethod("parseFrom", byte[].class);
            } catch (NoSuchMethodException e) {
                log.warn("No parseFrom method found for class: {}", cls.getName());
                return null;
            }
        });
        
        if (parseMethod == null) {
            throw new NoSuchMethodException("No parseFrom method found for " + targetClass.getName());
        }
        
        return parseMethod.invoke(null, body);
    }

    @Override
    public Message<?> toMessage(Object payload, MessageHeaders headers, Object conversionHint) {
        return toMessage(payload, headers);
    }

    @Override
    public Message<?> toMessage(Object payload, MessageHeaders headers) {
        try {
            byte[] body;
            String contentType;

            if (payload instanceof String) {
                body = ((String) payload).getBytes(StandardCharsets.UTF_8);
                contentType = CONTENT_TYPE_TEXT_PLAIN;
            } else if (payload instanceof byte[]) {
                body = (byte[]) payload;
                contentType = CONTENT_TYPE_OCTET_STREAM;
            } else if (payload instanceof MessageLite) {
                body = toProtobufByteArray((MessageLite) payload);
                contentType = CONTENT_TYPE_PROTOBUF;
            } else {
                body = objectMapper.writeValueAsBytes(payload);
                contentType = CONTENT_TYPE_JSON;
            }

            return MessageBuilder.withPayload(body)
                    .copyHeaders(headers)
                    .setHeaderIfAbsent("contentType", contentType)
                    .build();
        } catch (Exception e) {
            log.error("Failed to convert payload to message: {}", payload.getClass().getName(), e);
            throw new MessageConversionException("Failed to convert payload to message", e);
        }
    }
    
    /**
     * 将Protobuf消息转换为字节数组，使用缓存提高性能
     * 
     * @param message Protobuf消息
     * @return 字节数组
     * @throws NoSuchMethodException 找不到方法
     * @throws IllegalAccessException 方法访问异常
     * @throws InvocationTargetException 方法调用异常
     */
    private byte[] toProtobufByteArray(MessageLite message) 
            throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
        Class<? extends MessageLite> messageClass = message.getClass();
        Method toByteArrayMethod = protobufToByteArrayMethodCache.computeIfAbsent(messageClass, cls -> {
            try {
                return cls.getMethod("toByteArray");
            } catch (NoSuchMethodException e) {
                log.warn("No toByteArray method found for class: {}", cls.getName());
                return null;
            }
        });
        
        if (toByteArrayMethod == null) {
            throw new NoSuchMethodException("No toByteArray method found for " + messageClass.getName());
        }
        
        return (byte[]) toByteArrayMethod.invoke(message);
    }
}